"""
Cylindrical coordinate grid implementations.
"""

using LinearAlgebra
using StaticArrays

"""
    CylindricalSymGrid <: AbstractCylindricalGrid

Cylindrical coordinate grid with cylindrical symmetry (r,z).

# Fields
- `num_r::Int`: Number of radial points
- `num_z::Int`: Number of axial points  
- `r_bounds::Tuple{Float64,Float64}`: Radial bounds [r_min, r_max]
- `z_bounds::Tuple{Float64,Float64}`: Axial bounds [z_min, z_max]

# Examples
```julia
# Cylindrical grid: r ∈ [0,5], z ∈ [-2,2]
grid = CylindricalSymGrid(32, 64, (0.0, 5.0), (-2.0, 2.0))
```
"""
struct CylindricalSymGrid <: AbstractCylindricalGrid
    num_r::Int
    num_z::Int
    r_bounds::Tuple{Float64,Float64}
    z_bounds::Tuple{Float64,Float64}
    
    function CylindricalSymGrid(num_r::Int, num_z::Int, 
                               r_bounds::Tuple{Float64,Float64}, 
                               z_bounds::Tuple{Float64,Float64})
        num_r > 0 || throw(ArgumentError("Number of radial points must be positive"))
        num_z > 0 || throw(ArgumentError("Number of axial points must be positive"))
        r_bounds[1] >= 0 || throw(ArgumentError("Minimum radius must be non-negative"))
        r_bounds[1] < r_bounds[2] || throw(ArgumentError("Invalid radial bounds"))
        z_bounds[1] < z_bounds[2] || throw(ArgumentError("Invalid axial bounds"))
        new(num_r, num_z, r_bounds, z_bounds)
    end
end

# Interface implementation for CylindricalSymGrid
ndim(::CylindricalSymGrid) = 2  # (r,z) coordinates
Base.size(grid::CylindricalSymGrid) = (grid.num_r, grid.num_z)
extent(grid::CylindricalSymGrid) = (grid.r_bounds, grid.z_bounds)

function spacing(grid::CylindricalSymGrid)
    dr = (grid.r_bounds[2] - grid.r_bounds[1]) / (grid.num_r - 1)
    dz = (grid.z_bounds[2] - grid.z_bounds[1]) / (grid.num_z - 1)
    (dr, dz)
end

function coordinate(grid::CylindricalSymGrid, idx)
    r = grid.r_bounds[1] + (idx[1] - 1) * spacing(grid)[1]
    z = grid.z_bounds[1] + (idx[2] - 1) * spacing(grid)[2]
    SVector(r, z)
end

function index(grid::CylindricalSymGrid, coord)
    r, z = coord[1], coord[2]
    dr, dz = spacing(grid)
    idx_r = round(Int, (r - grid.r_bounds[1]) / dr + 1)
    idx_z = round(Int, (z - grid.z_bounds[1]) / dz + 1)
    (idx_r, idx_z)
end

# Coordinate transformations
"""
    cylindrical_to_cartesian(r, φ, z)

Convert cylindrical coordinates (r, φ, z) to Cartesian (x, y, z).
"""
function cylindrical_to_cartesian(r, φ, z)
    x = r * cos(φ)
    y = r * sin(φ)
    return SVector(x, y, z)
end

"""
    cartesian_to_cylindrical(x, y, z)

Convert Cartesian coordinates (x, y, z) to cylindrical (r, φ, z).
"""
function cartesian_to_cylindrical(x, y, z)
    r = sqrt(x^2 + y^2)
    φ = atan(y, x)
    return SVector(r, φ, z)
end

# Volume elements for integration
"""
    volume_element(grid::CylindricalSymGrid, r::Float64)

Volume element for cylindrical coordinates: 2π * r
"""
function volume_element(grid::CylindricalSymGrid, r::Float64)
    return 2π * r
end

"""
    integrate(grid::CylindricalSymGrid, values::AbstractMatrix)

Integrate values over cylindrical grid with proper volume weighting.
"""
function integrate(grid::CylindricalSymGrid, values::AbstractMatrix)
    dr, dz = spacing(grid)
    total = 0.0
    
    for i in 1:grid.num_r, j in 1:grid.num_z
        r = coordinate(grid, (i, j))[1]
        weight = volume_element(grid, r)
        total += values[i, j] * weight * dr * dz
    end
    
    return total
end

# Special methods for cylindrical symmetry
"""
    get_radial_profile(grid::CylindricalSymGrid, values::AbstractMatrix, z_idx::Int)

Extract radial profile at given z-index.
"""
function get_radial_profile(grid::CylindricalSymGrid, values::AbstractMatrix, z_idx::Int)
    return values[:, z_idx]
end

"""
    get_axial_profile(grid::CylindricalSymGrid, values::AbstractMatrix, r_idx::Int)

Extract axial profile at given r-index.
"""
function get_axial_profile(grid::CylindricalSymGrid, values::AbstractMatrix, r_idx::Int)
    return values[r_idx, :]
end